// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/base/exceptions.h>

namespace kumo::pollux {
    constexpr int kSortingNetworkMaxSize = 16;

    template<typename T, typename LessThan = std::less<T> >
    void sortingNetwork(T *data, int size, LessThan &&lt = {});

    namespace detail {
        // Compile time generated Bose-Nelson sorting network.
        //
        // https://bertdobbelaere.github.io/sorting_networks.html
        // https://github.com/Vectorized/Static-Sort/blob/master/include/static_sort.h
        template<int kSize>
        class SortingNetworkImpl {
        public:
            template<typename T, typename LessThan>
            static void apply(T *data, LessThan &&lt) {
                PS<T, LessThan, 1, kSize, (kSize <= 1)> ps(data, lt);
            }

        private:
            template<int I, int J, typename T, typename LessThan>
            static void compareExchange(T *data, LessThan lt) {
                // This is branchless if `lt' is branchless.
                auto c = lt(data[I], data[J]);
                auto min = c ? data[I] : data[J];
                data[J] = c ? data[J] : data[I];
                data[I] = min;
            }

            template<typename T, typename LessThan, int I, int J, int X, int Y>
            struct PB {
                PB(T *data, LessThan lt) {
                    enum {
                        L = X >> 1,
                        M = (X & 1 ? Y : Y + 1) >> 1,
                        IAddL = I + L,
                        XSubL = X - L,
                    };
                    PB<T, LessThan, I, J, L, M> p0(data, lt);
                    PB<T, LessThan, IAddL, J + M, XSubL, Y - M> p1(data, lt);
                    PB<T, LessThan, IAddL, J, XSubL, M> p2(data, lt);
                }
            };

            template<typename T, typename LessThan, int I, int J>
            struct PB<T, LessThan, I, J, 1, 1> {
                PB(T *data, LessThan lt) {
                    compareExchange<I - 1, J - 1>(data, lt);
                }
            };

            template<typename T, typename LessThan, int I, int J>
            struct PB<T, LessThan, I, J, 1, 2> {
                PB(T *data, LessThan lt) {
                    compareExchange<I - 1, J>(data, lt);
                    compareExchange<I - 1, J - 1>(data, lt);
                }
            };

            template<typename T, typename LessThan, int I, int J>
            struct PB<T, LessThan, I, J, 2, 1> {
                PB(T *data, LessThan lt) {
                    compareExchange<I - 1, J - 1>(data, lt);
                    compareExchange<I, J - 1>(data, lt);
                }
            };

            template<typename T, typename LessThan, int I, int M, bool kStop>
            struct PS {
                PS(T *data, LessThan lt) {
                    enum { L = M >> 1, IAddL = I + L, MSubL = M - L };
                    PS<T, LessThan, I, L, (L <= 1)> ps0(data, lt);
                    PS<T, LessThan, IAddL, MSubL, (MSubL <= 1)> ps1(data, lt);
                    PB<T, LessThan, I, IAddL, L, MSubL> pb(data, lt);
                }
            };

            template<typename T, typename LessThan, int I, int M>
            struct PS<T, LessThan, I, M, true> {
                PS(T * /*data*/, LessThan /*lt*/) {
                }
            };
        };
    } // namespace detail

    template<typename T, typename LessThan>
    void sortingNetwork(T *data, int size, LessThan &&lt) {
        switch (size) {
            case 0:
            case 1:
                return;

#ifdef POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE
#error "Macro name clash: POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE"
#endif
#define POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(_n)                            \
  case _n:                                                                   \
    detail::SortingNetworkImpl<_n>::apply(data, std::forward<LessThan>(lt)); \
    return;

            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(2)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(3)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(4)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(5)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(6)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(7)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(8)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(9)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(10)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(11)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(12)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(13)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(14)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(15)
            POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE(16)

#undef POLLUX_SORTING_NETWORK_IMPL_APPLY_CASE

            default:
                POLLUX_UNREACHABLE();
        }
    }
} // namespace kumo::pollux
